import sqlite3
import csv
import json
from typing import List, Dict, Optional, Union


class LocalDBManager:
    def __init__(self, db_path):
        """初始化数据库连接"""
        self.db_path = db_path
        self.conn = None
        self.cursor = None
        self.connect()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def connect(self):
        """建立数据库连接"""
        try:
            self.conn = sqlite3.connect(self.db_path)
            self.conn.row_factory = sqlite3.Row  # 使查询结果支持字典访问
            self.cursor = self.conn.cursor()
            print(f"成功连接到数据库: {self.db_path}")
        except sqlite3.Error as e:
            print(f"数据库连接失败: {e}")

    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.commit()
            self.conn.close()
            print("数据库连接已关闭")

    def execute_script(self, sql_file: str):
        """执行SQL脚本文件"""
        try:
            with open(sql_file, 'r', encoding='utf-8') as f:
                sql_script = f.read()
            self.cursor.executescript(sql_script)
            self.conn.commit()
            print(f"成功执行SQL脚本: {sql_file}")
        except Exception as e:
            print(f"执行SQL脚本失败: {e}")
            self.conn.rollback()

    def create_table(self, table_name: str, columns: Dict[str, str]):
        """创建数据表"""
        columns_str = ', '.join([f"{k} {v}" for k, v in columns.items()])
        sql = f"CREATE TABLE IF NOT EXISTS {table_name} ({columns_str})"
        self.execute(sql)

    def execute(self, sql: str, params: tuple = (), commit: bool = True):
        """执行SQL语句"""
        try:
            self.cursor.execute(sql, params)
            if commit:
                self.conn.commit()
            return True
        except sqlite3.Error as e:
            print(f"执行SQL失败: {e}\nSQL: {sql}")
            self.conn.rollback()
            return False

    def executemany(self, sql: str, params: tuple = (), commit: bool = True):
        """执行SQL语句"""
        try:
            self.cursor.executemany(sql, params)
            if commit:
                self.conn.commit()
            return True
        except sqlite3.Error as e:
            print(f"执行SQL失败: {e}\nSQL: {sql}")
            self.conn.rollback()
            return False

    def fetch(self, sql: str, params: tuple = (), fetch_all: bool = True):
        """执行查询语句"""
        try:
            self.cursor.execute(sql, params)
            return self.cursor.fetchall() if fetch_all else self.cursor.fetchone()
        except sqlite3.Error as e:
            print(f"查询失败: {e}\nSQL: {sql}")
            return None

    # ---------- CRUD 操作 ----------
    def insert(self, table: str, data: Dict) -> int:
        """插入单条数据"""
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['?'] * len(data))
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        if self.execute(sql, tuple(data.values())):
            return self.cursor.lastrowid
        return -1

    def batch_insert(self, table: str, columns: Union[List[str], str] = None,
                     data: Union[List[tuple], List[dict], tuple, dict] = None):
        """终极版批量插入方法
        参数:
            table: 表名
            columns: 可以是以下形式之一:
                - None: 自动推断所有字段(排除自增主键)
                - 列表: 指定要插入的字段
                - 字符串'*': 强制包含所有字段(包括自增主键)
            data: 可以是:
                - 单个元组: 插入单行数据
                - 单个字典: 插入单行数据
                - 元组列表: 每个元组对应一行数据
                - 字典列表: 每个字典对应一行数据
        """
        if data is None:
            print("警告：插入数据为空")
            return False

        try:
            # 标准化数据格式为列表
            if isinstance(data, (tuple, dict)):
                data = [data]

            # 获取完整表结构
            self.cursor.execute(f"PRAGMA table_info({table})")
            table_info = self.cursor.fetchall()
            all_columns = [col[1] for col in table_info]

            # 确定要插入的字段
            if columns == '*':
                insert_columns = all_columns
            elif columns is None:
                # 自动排除自增主键
                autoinc_pk = [col[1] for col in table_info
                              if col[5] == 1 and "INTEGER" in col[2].upper()]
                insert_columns = [col for col in all_columns
                                  if col not in autoinc_pk]
            else:
                insert_columns = columns

            # 准备数据和参数
            if isinstance(data[0], dict):
                # 字典数据转换为元组列表
                values = []
                for row in data:
                    values.append(tuple(row.get(col, None) for col in insert_columns))
                placeholders = ', '.join(['?'] * len(insert_columns))
            else:
                # 元组数据直接使用
                values = data
                placeholders = ', '.join(['?'] * len(values[0])) if values else ''

            # 验证数据维度
            expected_params = len(insert_columns)
            for i, row in enumerate(values):
                if len(row) != expected_params:
                    raise ValueError(
                        f"数据维度不匹配 (行 {i + 1})\n"
                        f"表字段 ({expected_params}个): {insert_columns}\n"
                        f"行数据 ({len(row)}个): {row[:10]}{'...' if len(row) > 10 else ''}"
                    )

            # 构建并执行SQL
            sql = f"INSERT INTO {table} ({', '.join(insert_columns)}) VALUES ({placeholders})"
            print(sql)
            print(values)
            self.cursor.executemany(sql, values)
            self.conn.commit()
            return True

        except sqlite3.Error as e:
            print(f"批量插入失败: {e}\nSQL: {sql}")
            self.conn.rollback()
            return False
        except Exception as e:
            print(f"数据处理失败: {str(e)[:200]}")
            return False

    def update(self, table: str, where: Dict, data: Dict) -> int:
        """更新数据"""
        set_clause = ', '.join([f"{k} = ?" for k in data.keys()])
        where_clause = ' AND '.join([f"{k} = ?" for k in where.keys()])
        sql = f"UPDATE {table} SET {set_clause} WHERE {where_clause}"
        params = tuple(data.values()) + tuple(where.values())
        if self.execute(sql, params):
            return self.cursor.rowcount
        return 0

    def delete(self, table: str, where: Dict) -> int:
        """删除数据"""
        where_clause = ' AND '.join([f"{k} = ?" for k in where.keys()])
        sql = f"DELETE FROM {table} WHERE {where_clause}"
        if self.execute(sql, tuple(where.values())):
            return self.cursor.rowcount
        return 0

    def query(self,
              table: str,
              columns: Union[str, List[str]] = "*",
              where: Optional[Dict[str, Union[str, int, float]]] = None,
              fetch_one: bool = False,
              order_by: Optional[str] = None,
              limit: Optional[int] = None) -> Union[List[dict], dict, None]:
        """
        增强版通用查询方法，支持操作符(如LIKE, >, <等)

        参数:
            where: 查询条件字典，如:
                  - 简单条件: {'name': '张三'} → name = '张三'
                  - 带操作符: {'age >': 20} → age > 20
                  - LIKE条件: {'name LIKE': '%张%'} → name LIKE '%张%'
        """
        sql = ""
        try:
            # 处理字段选择
            if isinstance(columns, str):
                columns = [columns] if columns != "*" else ["*"]
            columns_str = ", ".join(columns)

            # 构建WHERE子句
            where_clause = ""
            where_values = []
            if where:
                where_parts = []
                for key, value in where.items():
                    # 分离字段名和操作符
                    if " " in key:  # 包含操作符的情况
                        field, operator = key.rsplit(" ", 1)
                        where_parts.append(f"{field} {operator} ?")
                    else:  # 默认使用=操作符
                        where_parts.append(f"{key} = ?")
                    where_values.append(value)
                where_clause = " WHERE " + " AND ".join(where_parts)

            # 构建完整SQL
            sql = f"SELECT {columns_str} FROM {table}{where_clause}"
            if order_by:
                sql += f" ORDER BY {order_by}"
            if limit is not None:
                sql += f" LIMIT {limit}"

            # 执行查询
            self.cursor.execute(sql, where_values)

            # 获取结果
            if fetch_one:
                result = self.cursor.fetchone()
                return dict(result) if result else None
            else:
                return [dict(row) for row in self.cursor.fetchall()]

        except sqlite3.Error as e:
            print(f"查询失败: {e}\nSQL: {sql}")
            return None
        except Exception as e:
            print(f"查询处理失败: {e}")
            return None

    def get_count(self, table: str, where: Dict = None) -> int:
        """获取记录数量"""
        sql = f"SELECT COUNT(*) FROM {table}"
        params = ()
        if where:
            where_clause = ' AND '.join([f"{k} = ?" for k in where.keys()])
            sql += f" WHERE {where_clause}"
            params = tuple(where.values())
        result = self.fetch(sql, params, fetch_all=False)
        return result[0] if result else 0

    # ---------- 数据导入导出 ----------
    def export_to_csv(self, table: str, filename: str):
        """导出表数据到CSV文件"""
        try:
            data = self.query(table)
            with open(filename, 'w', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                # 写入列头
                writer.writerow([desc[0] for desc in self.cursor.description])
                # 写入数据
                for row in data:
                    writer.writerow(row)
            print(f"成功导出数据到 {filename}")
            return True
        except Exception as e:
            print(f"导出数据失败: {e}")
            return False

    def import_from_csv(self, table: str, filename: str):
        """从CSV文件导入数据"""
        try:
            with open(filename, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                columns = next(reader)  # 读取列头
                data = [tuple(row) for row in reader]
                return self.batch_insert(table, columns, data)
        except Exception as e:
            print(f"导入数据失败: {e}")
            return False

    def export_to_json(self, table: str, filename: str):
        """导出表数据到JSON文件"""
        try:
            data = self.query(table)
            result = [dict(row) for row in data]
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            print(f"成功导出数据到 {filename}")
            return True
        except Exception as e:
            print(f"导出数据失败: {e}")
            return False

    # ---------- 事务管理 ----------
    def begin_transaction(self):
        """开始事务"""
        self.execute("BEGIN TRANSACTION", commit=False)

    def commit(self):
        """提交事务"""
        self.conn.commit()

    def rollback(self):
        """回滚事务"""
        self.conn.rollback()
